
module HistTree where


import Data.Map
import Data.List
import Data.Either
import Data.Maybe
-- import Debug.Trace.Location
import Numeric

-- Left (b) == Node
-- Right    == map for more data ...
data HistTree a b = HTN {hmap :: (Map a (HistTree a b)),
			 node :: Maybe b }
    deriving (Show)

empty :: HistTree a b
empty = HTN {hmap = Data.Map.empty, node = Nothing}

insert :: (Ord a,Show a, Show b) => [a] -> b -> HistTree a b -> HistTree a b
insert [] v n0 = n0 {node = (Just v)}
insert (k:ks) v n0 =
    if  (isJust lookedup) -- (trace ("IN: " ++ (show (k:ks)) ++ "\n") (isJust lookedupeither))
       then n0 { hmap = (Data.Map.insert k (HistTree.insert ks v (fromJust lookedup)) (hmap n0))}
       else n0 { hmap = (Data.Map.insert k (HistTree.insert ks v HistTree.empty) (hmap n0))}
    where lookedup = Data.Map.lookup k (hmap n0)
          newnode = HistTree.empty

{-
lookupAverage :: (Ord a,Floating b) => [a] -> HistTree a b -> Maybe b
lookupAverage (t:ts) n = 
	let m = hmap n
	    val = Data.Map.lookup t m
	in if isJust val
	      then either (\x -> (Just x)) -- i.e. b value
	             (lookupAverage ts) -- i.e. HistTree value
	             (fromJust val)
	      else Just (htAverage n) -- can't find it
-}


lookup :: (Ord a) => [a] -> HistTree a b -> Maybe b
lookup [] n = node n
lookup (t:ts) n = 
	let m = hmap n
	    val = Data.Map.lookup t m
	in if isJust val
	      then HistTree.lookup ts (fromJust val)
	      else Nothing


derefTree :: (Ord a) => [a] -> HistTree a b -> HistTree a b
derefTree [] tree = tree
derefTree (t:ts) ht = derefTree ts (Data.Map.findWithDefault HistTree.empty t (hmap ht))

findWithDefault :: (Ord a) => b -> [a] -> HistTree a b -> b
findWithDefault def k t  = 
	if isJust (HistTree.lookup k t)
	   then fromJust $ HistTree.lookup k t
	   else def

update :: (Ord a,Show a,Show b) => (b -> Maybe b) -> [a] -> HistTree a b -> HistTree a b
update upfunc [] ht = 
    let val = node ht
	in if isJust val
	   then ht {node = upfunc (fromJust val)}
	   else ht
update upfunc (k:ks) ht =
    let mmap = hmap ht
	val = Data.Map.lookup k mmap
	in if isJust val
	      then ht {hmap = Data.Map.insert k (HistTree.update upfunc ks (fromJust val)) mmap}
	      else ht

-- updatefunc is a map from accval -> list of keys (where we are now, i.e. general key) -> value -> accval

accumulateNodes :: (HistTree b c) -> a -> (a -> [b] -> c -> a) -> a
accumulateNodes tree initval updatefunc = accumulateNodes1 tree initval [] updatefunc

accumulateNodes1 :: (HistTree b c) -> a -> [b] -> (a -> [b] -> c -> a) -> a
accumulateNodes1 tree initval cl updatefunc =
    let kvp = assocs $ hmap tree
	thisval = if (isJust (node tree)) 
                     then updatefunc initval cl (fromJust $ node tree) 
		     else initval
	in foldl (\x (y,z) -> accumulateNodes1 z x (y:cl) updatefunc) thisval kvp



{-
htAverage :: (Floating b) => HistTree a b -> b
htAverage n =
	let m = hmap n
	in if Data.Map.null m
	   then 0
	   else (foldl htacc 0 (elems m)) / (fromIntegral $ size m)

htacc :: (Floating b) => b -> (Either b (HistTree a b)) -> b
htacc x (Left y) = x + y
htacc x (Right m) = x + (htAverage m)
-}

-- showTree printlinefunc depth tree

showTreeParts :: (Show a,Ord a,Show b,Ord c,Show c,RealFloat c) => (Maybe b -> (c,String))  -> (a -> String) -> Int -> HistTree a b -> String
showTreeParts f printer depth m = 
    (foldl (\c r -> c ++ "---\n" ++ r) "" (Data.List.map partprint groups))
    where partprint part = foldl (++) "" (mtake $ reverse $ Data.List.map eleprint (sortBy wcomp part))
          wcomp (x1,_,l1) (x2,_,l2) = if x1 == x2
				         then l1 `compare` l2
					 else x1 `compare` x2
	  initdata = (cassocs f [] m)
	  groups = (groupBy lengthtest (sortBy sorter initdata))
	  sorter (w1,_,l1) (w2,_,l2) = (length l2) `compare` (length l1)
	  eleprint x = pprintelem f printer x
	  lengthtest (w1,_,l1) (w2,_,l2) = (length l1) == (length l2)
	  mtake = if depth /= 0
	             then take depth
		     else id

pprintelem f printer (w,no,list) = (snd $ f no) ++ " -> " ++  (foldl (\c r -> c ++ (tabTo (printer r) tabval )) "" list) ++ "\n" ++ (weightlist w list 10 0.6) ++ "\n"

weightlist w pos max lambda = foldl (++) "" [ (showFFloat (Just 2) (w * (lambda**(fromIntegral p))) " ") | p <- powlist ] 
    where powlist = (reverse [0 .. ((length pos) - 1)] ) ++ (reverse [-(max - (length pos)) .. -1])

cassocs :: (Show a,Show b,Ord c) =>  (Maybe b -> (c,String)) -> [a] -> HistTree a b -> [(c,Maybe b,[a])]
cassocs f ls m =
    let no = node m
	(w,_) = f no
	pairs = assocs $ hmap m
	thisitem = if isJust no then [(w,no,(reverse ls))] else []
    in (thisitem ++ (foldl (\st (k,m) -> st ++ (cassocs f (ls++[k]) m)) [] pairs))


-- showTree printlinefunc depth tree
showTree :: (Show a,Show b,Ord c) => (Maybe b -> (c,String)) -> (a -> String) -> Int -> HistTree a b -> String
showTree f printer depth m = foldl (++) "" (mtake $ reverse $ Data.List.map snd (sortBy wcomp (pprint f printer [] m 0)))
    where wcomp (x1,l1) (x2,l2) = if x1 == x2
				  then l1 `compare` l2
				  else x1 `compare` x2
	  mtake = if depth /= 0
	             then take depth
		     else id

pprint :: (Show b,Ord c) => (Maybe b -> (c,String)) -> (a -> String) -> [a] -> HistTree a b -> Int -> [(c,String)]
pprint f printer ls m tabcount =
    let no = node m
	pairs = assocs $ hmap m
	kp1 l = foldl (\a b -> a ++ (tabTo (printer b) tabval)) "" (reverse l) 
	pp1 map s = pprint f printer (ls++[s]) map (tabcount + 1)
	(w,nostr) = f no
	thisitem = if isJust no then (w,((tabTo (nostr ++ "  -->") tabval) ++ (kp1 ls) ++ "\n")) else (w,"")
    in (thisitem:(foldl (\st (k,m) -> st ++ (pp1 m k)) [] pairs))

tabTo str n = if length str > n - 1
                 then (take (n-1) str) ++ " "
                 else str ++ (replicate (n - (length str)) ' ')

tabval = 18 :: Int